# cython: profile=True
from v3dsolver.fluid_elements.fluid_elems import Saffman,FluidElem
from numpy import array,uint8,ones,empty
from v3dsolver.vector import Vec
try:
    from mpi4py import MPI
except ImportError:
    print "cannot import MPI, continuing.."
import pdb,time

class VelCalculator:
    """
    Object dealing with velocity calculations on a bunch of elements based on a
    tree data structure.
    """
    def __init__(self,tree=None):
        """
        Initialization of VelCalculator object. 
        
        tree: a TreeManager instance (keep default None if direct velocity
        computation is desired)
        """
        self.tree = tree 
    
    def direct_vel_computation(self,elems):
        """
        Updates the velocity of effect elems because of cause elems using direct
        calculation.
        In case some elements are both cause and effect, include them in both 
        cause_elems and effect_elems.
        """
        for effect in elems:
            if effect.is_effect():
                effect.vel = Vec()
                for cause in elems:
                    if cause.is_cause():
                        v = cause.get_vel(effect)
                        effect.vel.add_inplace(v)

    def serial_vel_computation(self):
        """
        Serially compute velocity on all effect elements based on the 
        generated tree.
        """
        self.tree.build_tree()
        for elem in self.tree.root.effect_elems:
            elem.vel = Vec()
        self.tree.generate_pseudo_particles()
        self.traverse_own_tree()
    
    def traverse_own_tree(self):
        """
        Traverses the tree made by current processor to calculate local 
        velocities for all its particles.
        """
        self.comm = MPI.COMM_WORLD
        root = self.tree.root
        for elem in self.tree.root.effect_elems:
            v = root.compute_vel(elem)
            elem.vel.add_inplace(v)
    
    def gather_celldata(self,level):
        """
        Gathers cell centers of cells at 'level' and returns the final 
        concatenated list.
        """
        celldata = [self.comm.Get_rank(),[]]
        for lev in range(level): # for cells at level less than 'level'
            for cell in self.tree.cells[lev]:
                if cell.childless and not cell.effect_elems==[]:
                    # add 'em if childless and non-empty
                    celldata[1] += [(cell.center,cell.id)] 
        for cell in self.tree.cells[level]:  # cells at 'level'
            celldata[1] += [(cell.center,cell.id)]
        all_celldata = self.comm.allgather(sendobj=celldata)
        return all_celldata
    
    def mark_nearby_faroff(self,cell,cellcenter,level):
        """
        Marks 'cell' a nearby/far-off cell to the cell defined by cellcenter at 
        level 'level'.
        
        Returns a list of nearby cells and a list of far-off cells to 
        cellcenter.
        """
        nearby = []
        faroff = []
        maxlevel = level - 1 # level is level of cellcenter cell
        if not cell.is_neighbour(cellcenter,level): # far-off
            faroff += [cell]
        else:                                       # near-by
            if cell.level<maxlevel:
                if cell.childless and not cell.cause_elems==[]: # childless and non-empty
                    nearby += [cell]
                else:
                    for child in cell.children:
                        childnearby, childfaroff = self.mark_nearby_faroff(child,
                                                              cellcenter,level)
                        nearby += childnearby
                        faroff += childfaroff
            elif cell.level==maxlevel:
                nearby += [cell]
            else:
                pass
        return nearby, faroff
    
    def gather_nearby_faroff_data(self,all_celldata,level):
        """
        Track neighbours and far-off cells
        """
        myrank = self.comm.rank
        send_objects = [myrank] # objects sent by myrank processor
        for i in range(len(all_celldata)):
            destrank = all_celldata[i][0]
            if not destrank==myrank:
                cells_data = []
                for j in range(len(all_celldata[i][1])):
                    celldata = all_celldata[i][1][j]
                    cellcenter = celldata[0]
                    cellid = celldata[1]
                    nearby,faroff = self.mark_nearby_faroff(self.tree.root,
                                                            cellcenter,level)
                    nearby_ids = [cell.id for cell in nearby]
                    pseudo_pos_vors = []
                    for cell in faroff:
                       for pseudo in cell.pseudo_p:
                            pseudo_pos_vors += [[pseudo.pos,pseudo.vor]]
                    cells_data += [[cellid,nearby_ids,pseudo_pos_vors]]
                sendobj = [destrank,cells_data] 
                send_objects += [sendobj]
        all_nearby_faroff_data = self.comm.allgather(send_objects) 
        return all_nearby_faroff_data
    
    def compute_faroff_vel(self,data):    
        """
        Computes velocity due to pseudo particles of far-off cells.
        """
        for i in range(len(data)):
            celldata = data[i]
            cellid = celldata[0] 
            pseudo_pos_vors = celldata[2]
            cell = self.tree.get_cell(cellid)
            for pseudo_pos_vor in pseudo_pos_vors:
                pseudop = Saffman(pseudo_pos_vor[0],pseudo_pos_vor[1],
                                  vel=Vec(), rad=self.tree.rad)
                for elem in cell.effect_elems:
                    # calculating velocities due to far-off cells
                    v = pseudop.get_vel(elem)
                    elem.vel.add_inplace(v) 
    
    def send_particles(self,data,server_rank):
        """
        Sends particles to processor having rank equal to server_rank for 
        nearby velocity computation. The nearby_ids etc. are stored in data 
        in such a manner that for every cell, it has the following form:

        [cellid,nearby_ids,pseudo_pos_vor]
        """
        sendback_particles = []
        for i in range(len(data)):
            celldata = data[i]
            cellid = celldata[0]
            nearby_ids = celldata[1]
            cell = self.tree.get_cell(cellid)
            elem_positions = [elem.pos for elem in cell.effect_elems]
            if not elem_positions is []:
                sendback_particles.append([cellid,nearby_ids,elem_positions])
        return sendback_particles
    
    def vel_recvd_particles(self,recvd_data,client_rank):
        """
        Based on the received particle data from the processor having rank client_rank, it 
        calculates velocities on them and sends back.
        """        
        send_vels_obj = []
        for j in range(len(recvd_data)):
            celldata = recvd_data[j]
            cellid = celldata[0]
            nearby_ids = celldata[1]
            elem_positions = celldata[2]
            nearby_cells = [self.tree.get_cell(id) for id in nearby_ids]
            vels = [Vec() for i in range(len(elem_positions))]
            for k in range(len(elem_positions)):
                elem_pos = elem_positions[k]
                dummy_elem = FluidElem(pos=elem_pos,vel=Vec(),cause_or_effect=1)
                for cause_cell in nearby_cells:
                    v = cause_cell.compute_vel(dummy_elem)
                    vels[k].add_inplace(v)
            send_vels_obj.append([cellid,vels])
        return send_vels_obj
    
    def assign_nearby_vels(self,recvd_vel_data):
        """
        Assigns nearby velocities to all the particles based on obtained 
        nearby velocity data from a particular processor.
        """
        for j in range(len(recvd_vel_data)):
            celldata = recvd_vel_data[j]
            cellid = celldata[0]
            velocities = celldata[1]
            cell = self.tree.get_cell(cellid)
            for i in range(len(cell.effect_elems)):
                elem = cell.effect_elems[i]
                elem.vel.add_inplace(velocities[i])

    def share_data(self,data):
        """
        Shares a particular data among all processors using blocking
        send/recv operations. 

        Requires:
        data: a list containing data to be sent to each processor, with
              the list being indexed according to processor ranks
        Returns:
        recvd_data: a list containing data obtained from different processors,
                    with the list being indexed according to processor ranks
        """
        recvd_data = [None for i in range(self.comm.size)]
        
        # first half of processors sending, second half receiving
        for i in range(self.comm.size):
            if i>self.comm.rank:
                self.comm.send(data[i],dest=i,tag=0)
            if i==self.comm.rank:
                pass
            if i<self.comm.rank:
                recvd_data[i] = self.comm.recv(source=i,tag=0)
        # second half of processors sending, first half receiving
        for i in range(self.comm.size):
            if i<self.comm.rank:
                self.comm.send(data[i],dest=i,tag=1)
            if i==self.comm.rank:
                pass
            if i>self.comm.rank:
                recvd_data[i] = self.comm.recv(source=i,tag=1)
         
        return recvd_data
    
    def parallel_vel_computation(self):
        """
        Parallel computation of velocity. Steps involved are (attributes 
        written in leading square brackets denote the role current processor 
        takes while performing that step):

        1.[Client] Traverse own tree data strucuture. This part doesn't involve 
                   any non-local communication.
        2.[Server] AllGather cell centers at 'level'
        3.[Server] Based on this cell center data, determine nearby/faroff 
                   cells to different cells. If far-off, mark its pseudo 
                   particles. If nearby, do same check for its children (go 
                   this way only level+1 deep). At maxlevel (or the state of 
                   childlessness), mark the cell as nearby. Send this faroff 
                   (pseudoparticles) and nearby data to client.
        4.[Client] Based on this nearby/far-off data, compute vel due to faroff 
                   cells using obtained pseudo particles. For nearby cells, 
                   send particle positions to these nearby cells for velocity 
                   computation.
        5.[Server] Compute vels at obtained particle positions by traversing 
                   local tree. Send these velocities back.
        6.[Client] Update particle velocities based on obtained nearby 
                   contribution obtained from server.
        """
        # Preprocessing for velocity calculation
        self.comm = self.tree.comm
        self.tree.build_tree()
        for elem in self.tree.root.effect_elems:
            elem.vel = Vec()
        self.tree.generate_pseudo_particles()
       
        # Every process traverses its own tree data structure to find local 
        # velocity contribution
        self.traverse_own_tree() 	
        
        # Servers gather cellcenter data at given level (default: 3)
        level = min(self.tree.maxlevel,3)
        all_celldata = self.gather_celldata(level) 

        # Clients gather nearby/faroff data generated by server
        all_nearby_faroff_data = self.gather_nearby_faroff_data(all_celldata,
                                                                level) 
        particle_data = [None for i in range(self.comm.size)]
        for nearfar_data in all_nearby_faroff_data:
            server = nearfar_data[0]
            for src_data in nearfar_data[1:]:
                client = src_data[0]
                if self.comm.Get_rank()==client:
                    self.compute_faroff_vel(src_data[1]) 
                    particle_data[server] = self.send_particles(src_data[1],
                                                                server)
        
        # All particle data is shared between processors                                                        
        recvd_particle_data = self.share_data(particle_data)
        
        # Server computes velocities at received particles (and sends back 
        # computed velocities)
        vel_data = [None for i in range(self.comm.size)]
        for i in range(len(recvd_particle_data)):
            data = recvd_particle_data[i]
            if not data is None:
                vel_data[i] = self.vel_recvd_particles(data,i)
        
        # All velocity data is shared between processors
        recvd_vel_data = self.share_data(vel_data)
        
        # Client updates particle velocities according to the received velocity data
        for data in recvd_vel_data:
            if not data==None:
                self.assign_nearby_vels(data)
